package com.xyq.judge.config;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.io.FileUtil;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.*;
import com.github.dockerjava.api.model.*;
import com.xyq.common.core.constants.JudgeConstants;

import java.io.File;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;

// docker容器池
public class DockerSandBoxPool {

    private DockerClient dockerClient;
    private  String sandboxImage;
    private String volumeDir;
    private   Long memoryLimit;

    private Long memorySwapLimit;
    private Long cpuLimit;
    private int poolSize;
    private String containerNamePrefix;

    // 阻塞队列,存储各个容器id
    private BlockingQueue<String> containerQueue ;

    // 通过一个Map 来存储 容器id和名字的映射
    private Map<String , String> containerNameMap ;

    public DockerSandBoxPool(DockerClient dockerClient, String sandboxImage, String volumeDir, Long memoryLimit,
                             Long memorySwapLimit, Long cpuLimit, int poolSize, String containerNamePrefix) {
        this.dockerClient = dockerClient;
        this.sandboxImage = sandboxImage;
        this.volumeDir = volumeDir;
        this.memoryLimit = memoryLimit;
        this.memorySwapLimit = memorySwapLimit;
        this.cpuLimit = cpuLimit;
        this.poolSize = poolSize;
        this.containerNamePrefix = containerNamePrefix;
        // 构造阻塞队列
        this.containerQueue = new ArrayBlockingQueue<>(poolSize);
        // 容器名字和id映射
        this.containerNameMap = new HashMap<>() ;
    }

    // 初始化容器池
    public void initDockerPool(int poolSize){
        // 创建容器
        for (int i = 0; i < poolSize; i++) {
            // 每个容器的名字: 前缀 + 序号
            createContainer(containerNamePrefix +"-" + i) ;
        }
    }

    private void createContainer(String containerName) {

        // 创建容器之前,先判断容器池里是否已经有这个容器了
        // 有就不用创建了

        // 获取所有容器列表 , 判断容器列表里是否已经有了这个容器
        List<Container> containerList = dockerClient.listContainersCmd().withShowAll(true).exec();
        if(!CollectionUtil.isEmpty(containerList)){
            // docker获取的名字是 有 / 前缀的
          String names = JudgeConstants.JAVA_CONTAINER_PREFIX + containerName ;
           for(Container container : containerList){
               String[] containerNames = container.getNames() ;
               if(containerNames != null && containerNames.length>0 && names.equals(containerNames[0])){
                   if("created".equals(container.getState()) || "exited".equals(container.getState())){
                       // 如果容器未启动, 启动容器
                       dockerClient.startContainerCmd(container.getId()).exec() ;
                   }
                  // 已经有容器了 ,就直接加到 容器池里 , 不需要重新创建
                   containerQueue.add(container.getId());
                   // 添加一个容器id和 名字的映射
                   containerNameMap.put(container.getId(),  containerName) ;
                   return;
               }
           }
        }


        //拉取镜像: 需要一个java8的镜像
        pullJavaEnvImage();
        //创建容器  限制资源,权限
        HostConfig hostConfig = getHostConfig( containerName);
        // 创建容器命令
        CreateContainerCmd containerCmd = dockerClient
                // 使用哪个镜像
                .createContainerCmd(JudgeConstants.JAVA_ENV_IMAGE)
                // 容器的名字
                .withName(containerName);
        CreateContainerResponse createContainerResponse = containerCmd
                .withHostConfig(hostConfig)
                .withAttachStderr(true)   // 执行的结果无论是正确还是错误
                .withAttachStdout(true)   // 都英爱获取到
                .withTty(true)          // 配置一个伪终端,方便后续执行代码
                .exec();
        //记录容器id
        String containerId = createContainerResponse.getId();
        //启动容器
        dockerClient.startContainerCmd(containerId).exec();
        // 启动容器后 添加容器到 阻塞队列
        containerQueue.add(containerId) ;
        // 添加一个容器id和 名字的映射
        containerNameMap.put(containerId,  containerName) ;
    }

    // 从 容器池里获取一个容器
    public String getContainer(){
        try {
            return containerQueue.take() ;
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
   // 归还容器, 使用玩容器后归还容器给容器池
    public void returnContainer(String containerId){
        containerQueue.add(containerId);
    }


    //限制资源 ,容器的权限
    private HostConfig getHostConfig(String containerName) {
        HostConfig hostConfig = new HostConfig();
        //设置挂载目录，指定用户代码路径,

        // 挂载的目录 有多个容器 , 不能用同一个挂载目录
        // 给每个目录都配置一个自己的目录
        String userCodeDir = createContainerDir(containerName);
        hostConfig.setBinds(new Bind(userCodeDir, new Volume(volumeDir)));
        //限制docker容器使用资源
        hostConfig.withMemory(memoryLimit); // 内存限制
        hostConfig.withMemorySwap(memorySwapLimit); // 交换限制
        hostConfig.withCpuCount(cpuLimit);  // cpu限制
        hostConfig.withNetworkMode("none");  //禁用网络
        hostConfig.withReadonlyRootfs(true); //禁止在root目录写文件
        return hostConfig;
    }


    //拉取java执行环境镜像 需要控制只拉取一次
    private void pullJavaEnvImage() {
        ListImagesCmd listImagesCmd = dockerClient.listImagesCmd();
        List<Image> imageList = listImagesCmd.exec();
        // 遍历 docker 拥有的 镜像列表,如果有了java8的镜像就不需要再拉取镜像
        for (Image image : imageList) {
            String[] repoTags = image.getRepoTags();
            if (repoTags != null && repoTags.length > 0 && sandboxImage.equals(repoTags[0])) {
                return;
            }
        }
        PullImageCmd pullImageCmd = dockerClient.pullImageCmd(sandboxImage);
        try {
            //            用来监听 docker , 同时 .await 等待它拉取镜像完成
            pullImageCmd.exec(new PullImageResultCallback()).awaitCompletion();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }


    // 给每个容器指定 挂载文件 , 通过容器名字 ,生成 挂载目录
    private String createContainerDir(String containerName){
        // 一级目录 , 存放容器挂载目录
        String codeDir = System.getProperty("user.dir") + File.separator + JudgeConstants.CODE_DIR_POOL;
        if (!FileUtil.exist(codeDir)) {
            // 创建存放用户代码的目录
            FileUtil.mkdir(codeDir);
        }
        // 每个容器挂载的目录
      return  codeDir + File.separator + containerName ;
    }

     // 获取容器的挂载路径
    public String getCodeDir(String containerId) {
        String containerName = containerNameMap.get(containerId);
        return System.getProperty("user.dir") + File.separator + JudgeConstants.CODE_DIR_POOL +
                File.separator + containerName ;
    }
}
